#!/usr/bin/env python3
"""
树莓派机器人语音增强版 
功能：
1. 保留原有机器人控制功能
2. 将Coze AI回复的文本通过TTS转换为语音播放
3. 使用回调机制确保语音完全播放完成后再恢复语音识别
4. 彻底解决自监听问题
5. 使用cmn-latn-pinyin语音ID确保中文发音正确
"""

# 需要安装的依赖:
# sudo apt update
# sudo apt install espeak-ng espeak-ng-data
# pip3 install pyttsx3 spidev pyaudio vosk cozepy

import spidev
import time
import sys
import threading
import json
import pyaudio
import os
import termios
import tty
import select
import pyttsx3  # 新增TTS库
from vosk import Model, KaldiRecognizer
# Coze相关导入
from cozepy import COZE_CN_BASE_URL, TokenAuth, Coze
from cozepy import Message, ChatStatus, MessageContentType, ChatEventType

# 命令定义
# CMD_IDLE    = 0x00
# CMD_STAND   = 0x01
# CMD_WALK_F  = 0x02
# CMD_WALK_B  = 0x03
# CMD_LIE     = 0x04
# CMD_SIT     = 0x05
# CMD_STATUS  = 0x55

CMD_IDLE    =    0x00
CMD_STAND   =    0x01
CMD_WALK_FWD=    0x02
CMD_WALK_BWD=    0x03
CMD_RUN     =    0x04
CMD_TURN_L  =    0x05
CMD_TURN_R  =    0x06
CMD_LIE     =    0x07
CMD_SIT     =    0x08
CMD_STATUS  =    0x55


class CozeChat:
    def __init__(self, coze_api_token: str, bot_id: str, user_id: str, coze_api_base: str = COZE_CN_BASE_URL):
        """
        初始化CozeChat客户端
        :param coze_api_token: API访问令牌
        :param bot_id: Bot的ID
        :param user_id: 用户ID（可使用业务ID或随机字符串）
        :param coze_api_base: API基础URL（默认为COZE_CN_BASE_URL）
        """
        self.coze = Coze(
            auth=TokenAuth(token=coze_api_token),
            base_url=coze_api_base
        )
        self.bot_id = bot_id
        self.user_id = user_id
        
    def input_text(self, text: str):
        """
        发送用户输入文本并启动流式对话
        :param text: 用户输入的文本内容
        :return: 流式事件迭代器
        """
        return self.coze.chat.stream(
            bot_id=self.bot_id,
            user_id=self.user_id,
            additional_messages=[
                Message.build_user_question_text(text),
            ],
        )
        
    def output_text(self, event_stream):
        """
        处理流式事件并输出文本
        :param event_stream: 流式事件迭代器
        """
        full_response = []
        for event in event_stream:
            if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                content = event.message.content
                print(content, end="", flush=True)
                full_response.append(content)
            elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                print("\n")  # 完成时换行
                print(f"Token usage: {event.chat.usage.token_count}")
                print(f"Full response length: {len(''.join(full_response))} characters")
                return ''.join(full_response)
        # 如果没有收到完成事件（理论上不应该发生）
        return ''.join(full_response)
        
    def chat(self, user_message: str) -> str:
        """
        完整的聊天流程：发送消息并处理响应
        :param user_message: 用户输入的消息
        :return: 完整的响应文本
        """
        event_stream = self.input_text(user_message)
        return self.output_text(event_stream)

class VoiceRecognizer:
    def __init__(self, model_path="/home/jiang/Desktop/vosk-model-small-cn-0.22"):
        self.model_path = model_path
        self.lock = threading.Lock()
        self._reset_audio_resources()
        # 语音命令到动作的映射
        self.command_map = {
            "站立": "stand",
            "站起来": "stand",
            "站": "stand",
            "前进": "forward",
            "向前": "forward",
            "前": "forward",
            "后退": "backward",
            "向后": "backward",
            "向后走": "backward",
            "后": "backward",
            "奔跑": "run",
            "跑": "run",
            "向左转": "turn_left",
            "左转": "turn_left",
            "向左": "turn_left",
            "向右转": "turn_right",
            "右转": "turn_right",
            "向右": "turn_right",
            "趴下": "lie",
            "卧": "lie",
            "躺下": "lie",
            "趴": "lie",
            "坐下": "sit",
            "坐": "sit",
            "停止": "stop",
            "停": "stop",
            "状态": "status",
            "空闲": "stop",
            "聊天": "coze",
            "对话": "coze",
            "问问": "coze",
            "问": "coze"
        }
        # 控制变量
        self.active = False
        self.latest_command = ""
        self.coze_mode = False  # 是否处于Coze对话模式
        self.ignore_next_audio = False  # 新增：忽略下一次音频处理
        self.tts_playing = False  # 新增：跟踪TTS是否正在播放
        
    def _reset_audio_resources(self):
        """重置音频资源，确保可以重新初始化"""
        self.model = None
        self.recognizer = None
        self.p = None
        self.stream = None
        
    def _init_audio_resources(self):
        """初始化音频资源"""
        try:
            # 加载中文语音模型
            self.model = Model(self.model_path)
            self.recognizer = KaldiRecognizer(self.model, 16000)
            self.recognizer.SetWords(True)
            
            # 初始化音频输入
            self.p = pyaudio.PyAudio()
            self.stream = self.p.open(
                format=pyaudio.paInt16,
                channels=1,
                rate=16000,
                input=True,
                input_device_index=1,  # 根据实际设备调整
                frames_per_buffer=4096
            )
            return True
        except Exception as e:
            print(f"✗ 音频资源初始化失败: {str(e)}")
            self._reset_audio_resources()
            return False
            
    def start(self):
        """启动语音识别线程"""
        with self.lock:
            if self.active:
                return  # 如果已经启动，不再重复启动
                
            # 确保音频资源已初始化
            if not self._init_audio_resources():
                print("✗ 无法初始化语音识别，请检查音频设备和模型路径")
                return
                
            self.active = True
            self.thread = threading.Thread(target=self._recognize_loop)
            self.thread.daemon = True
            self.thread.start()
        print("✓ 语音识别已启动...")
        
    def stop(self):
        """停止语音识别"""
        with self.lock:
            if not self.active:
                return  # 如果已经停止，不再重复停止
            self.active = False
        
        # 等待线程结束
        if hasattr(self, 'thread') and self.thread.is_alive():
            self.thread.join(timeout=1.0)
            
        # 关闭音频流
        try:
            if self.stream:
                self.stream.stop_stream()
                self.stream.close()
        except Exception as e:
            print(f"⚠ 关闭音频流时出错: {str(e)}")
            
        try:
            if self.p:
                self.p.terminate()
        except Exception as e:
            print(f"⚠ 终止PyAudio时出错: {str(e)}")
            
        # 重置资源状态
        self._reset_audio_resources()
        print("✓ 语音识别已停止")
        
    def is_active(self):
        """检查语音识别是否处于活动状态"""
        with self.lock:
            return self.active
            
    def _clear_audio_buffer(self):
        """清除音频缓冲区，防止残留音频被识别"""
        try:
            # 读取并丢弃缓冲区中的所有数据
            if self.stream and self.active:
                for _ in range(10):  # 读取10次缓冲区
                    self.stream.read(4096, exception_on_overflow=False)
            return True
        except Exception as e:
            print(f"✗ 清除音频缓冲区失败: {str(e)}")
            return False
            
    def _recognize_loop(self):
        """语音识别主循环"""
        while self.active:
            try:
                data = self.stream.read(4096, exception_on_overflow=False)
                
                # 新增：检查是否应忽略此音频数据
                with self.lock:
                    if self.ignore_next_audio:
                        self.ignore_next_audio = False
                        continue  # 忽略此音频数据
                
                # 新增：如果TTS正在播放，跳过识别
                with self.lock:
                    if self.tts_playing:
                        continue
                
                if self.recognizer.AcceptWaveform(data):
                    result = json.loads(self.recognizer.Result())
                    text = result.get('text', '').strip()
                    if text:
                        print(f"\n识别到语音: {text}")
                        # 检查是否是Coze触发命令
                        if not self.coze_mode:
                            for keyword, command in self.command_map.items():
                                if keyword in text:
                                    if command == "coze":
                                        self.coze_mode = True
                                        print("✓ 已进入Coze对话模式，请说话...")
                                        self.latest_command = "coze_start"
                                    else:
                                        self.latest_command = command
                                    break
                        else:
                            # 在Coze对话模式下，将所有语音作为输入
                            self.latest_command = f"coze:{text}"
                            self.coze_mode = False  # 重置模式，每次只处理一条Coze消息
                            print(f"Coze模式: 发送消息 '{text}'")
                # 显示部分识别结果
                else:
                    partial = json.loads(self.recognizer.PartialResult())
                    if 'partial' in partial:
                        sys.stdout.write('\r' + ' ' * 50)  # 清除当前行
                        sys.stdout.write(f"\r监听中: {partial['partial']}")
                        sys.stdout.flush()
            except Exception as e:
                # 重要修复：当音频资源可能已关闭时，检查状态
                with self.lock:
                    if not self.active:
                        break  # 如果已停止，退出循环
                print(f"\n✗ 语音识别错误: {str(e)}")
                time.sleep(1)
                
    def get_command(self):
        """获取最新的语音命令并清除"""
        with self.lock:
            cmd = self.latest_command
            self.latest_command = ""
            return cmd
            
    def pause_for_tts(self):
        """暂停语音识别并准备TTS播放（清除缓冲区）"""
        with self.lock:
            if not self.active:
                return False
                
            # 暂停识别
            self.active = False
            # 标记TTS正在播放
            self.tts_playing = True
            
        # 等待线程结束
        if hasattr(self, 'thread') and self.thread.is_alive():
            self.thread.join(timeout=1.0)
            
        # 清除音频缓冲区
        self._clear_audio_buffer()
        
        print("✓ 暂停语音识别并清除音频缓冲（防止自监听）")
        return True
        
    def resume_after_tts(self):
        """在TTS播放完成后恢复语音识别"""
        with self.lock:
            # 标记TTS已结束
            self.tts_playing = False
            
        # 额外延迟，确保TTS声音完全消失
        time.sleep(0.5)
        
        # 再次清除缓冲区
        self._clear_audio_buffer()
        
        # 标记下一次音频处理应被忽略（双重保护）
        with self.lock:
            self.ignore_next_audio = True
            
        # 重新启动
        self.start()
        print("✓ 恢复语音识别（TTS播放完成）")
        return True

class RobotController:
    def __init__(self, bus=0, device=0, speed=1000000):
        # 初始化SPI连接
        self.spi = spidev.SpiDev()
        try:
            self.spi.open(bus, device)
            self.spi.max_speed_hz = speed
            self.spi.mode = 0b00
            print(f"✓ SPI连接成功 (总线 {bus}, 设备 {device}, 速度 {speed}Hz)")
        except Exception as e:
            print(f"✗ SPI连接失败: {str(e)}")
            self.spi = None
            
        self.current_state = "未知"
        
        # 初始化TTS引擎
        try:
            print("✓ 正在初始化TTS引擎...")
            self.tts_engine = pyttsx3.init()
            # 设置正确的中文语音ID (cmn-latn-pinyin)
            self.tts_engine.setProperty('voice', 'cmn-latn-pinyin')
            # 优化语音参数
            self.tts_engine.setProperty('rate', 200)  # 语速
            self.tts_engine.setProperty('volume', 0.6)  # 音量
            print("✓ TTS引擎初始化成功，语音ID: cmn-latn-pinyin")
        except Exception as e:
            print(f"✗ TTS引擎初始化失败: {str(e)}")
            print("提示: 请确保已安装espeak-ng和中文语音包")
            print("      sudo apt install espeak-ng espeak-ng-data")
            self.tts_engine = None
            
        # 初始化语音识别
        self.voice_recognizer = VoiceRecognizer()
        self.voice_recognizer.start()
        
        # 初始化CozeChat（替换为你的实际参数）
        self.coze_client = CozeChat(
            coze_api_token="pat_i3sIq4ITqvvda3uNOpJ43cF1HssI3EbxCo7coECVBlskcz5CGA5s6sGJmaKD4Mm2",
            bot_id="7534957136225878055",
            user_id="robot_user_" + str(os.getpid())
        )
        print("✓ CozeChat客户端已初始化")
        
    def send_command(self, command):
        """通过SPI发送命令并返回Arduino响应"""
        if not self.spi:
            print("✗ SPI未连接，无法发送命令")
            return 0
        try:
            # SPI传输命令
            response = self.spi.xfer2([command])
            return response[0]
        except Exception as e:
            print(f"✗ SPI传输错误: {str(e)}")
            return 0
            
    def get_status(self):
        """通过SPI获取当前状态"""
        status_code = self.send_command(CMD_STATUS)
        self._update_state(status_code)
        return self.current_state
        
    def _update_state(self, status_code):
        """更新状态显示"""
        states = {
            CMD_IDLE: "空闲",
            CMD_STAND: "站立",
            CMD_WALK_FWD: "向前走",
            CMD_WALK_BWD: "向后走",
            CMD_RUN: "奔跑",
            CMD_TURN_L: "向左转",
            CMD_TURN_R: "向右转",
            CMD_LIE: "趴下",
            CMD_SIT: "坐下"

        }
        self.current_state = states.get(status_code, f"未知状态({status_code})")
        
    def stand(self):
        """通过SPI发送站立命令"""
        self.send_command(CMD_STAND)
        print("✓ 命令: 站立")
        
    def walk_forward(self):
        """通过SPI发送前进命令"""
        self.send_command(CMD_WALK_FWD)
        print("✓ 命令: 向前走")
        
    def walk_backward(self):
        """通过SPI发送后退命令"""
        self.send_command(CMD_WALK_BWD)
        print("✓ 命令: 向后走")

    def run(self):
        """通过SPI发送奔跑命令"""
        self.send_command(CMD_RUN)
        print("✓ 命令: 奔跑")

    def turn_left(self):
        """通过SPI发送向左转命令"""
        self.send_command(CMD_TURN_L)
        print("✓ 命令: 向左转")

    def turn_right(self):
        """通过SPI发送向右转命令"""
        self.send_command(CMD_TURN_R)
        print("✓ 命令: 向右转")

    
    def lie_down(self):
        """通过SPI发送趴下命令"""
        self.send_command(CMD_LIE)
        print("✓ 命令: 趴下")
        
    def sit(self):
        """通过SPI发送坐下命令"""
        self.send_command(CMD_SIT)
        print("✓ 命令: 坐下")
        
    def idle(self):
        """通过SPI发送停止命令"""
        self.send_command(CMD_IDLE)
        print("✓ 命令: 停止/空闲")
        
    def pause_voice_recognition(self):
        """暂停语音识别并清除音频缓冲（线程安全）"""
        return self.voice_recognizer.pause_for_tts()
        
    def resume_voice_recognition(self):
        """在TTS播放完成后恢复语音识别（线程安全）"""
        return self.voice_recognizer.resume_after_tts()
        
    def _tts_end_callback(self, name, completed):
        """TTS播放完成的回调函数"""
        if completed:
            print("✓ TTS播放完成，恢复语音识别")
            self.resume_voice_recognition()
        else:
            print("✗ TTS播放被中断")
            self.resume_voice_recognition()
            
    def text_to_speech(self, text):
        """将文本转换为语音并播放，同时暂停语音识别防止自监听"""
        if not self.tts_engine:
            print("✗ TTS引擎未初始化，无法播放语音")
            return False
            
        was_recognizing = False
        try:
            # 暂停语音识别并清除缓冲区，防止将TTS语音误认为用户输入
            was_recognizing = self.pause_voice_recognition()
            
            print(f"✓ 正在播放语音: {text[:30]}{'...' if len(text) > 30 else ''}")
            
            # 使用回调机制确保语音完全播放完成
            self.tts_engine.connect('finished-utterance', self._tts_end_callback)
            self.tts_engine.say(text)
            self.tts_engine.runAndWait()
            
            # 注意：runAndWait() 会阻塞直到语音完成
            # 但我们已经设置了回调，所以不需要额外等待
            
            return True
        except Exception as e:
            print(f"✗ 语音播放失败: {str(e)}")
            # 确保即使出错也恢复语音识别
            if was_recognizing:
                self.resume_voice_recognition()
            return False
            
    def process_coze_command(self, text):
        """处理Coze命令，将语音转为文本并发送给Coze"""
        print(f"\n[Coze] 用户输入: {text}")
        print("[Coze] 正在与AI对话...")
        try:
            # 使用方式一调用CozeChat
            response = self.coze_client.chat(text)
            print(f"\n[Coze] AI响应: {response[:100]}{'...' if len(response) > 100 else ''}")
            
            # 将Coze的回复转换为语音
            print("[Coze] 正在播放AI回复...")
            self.text_to_speech(response)
            
            # 这里可以根据Coze的响应执行特定操作
            # 例如，如果AI响应中包含"站立"，则让机器人站立
            if "站立" in response:
                self.stand()
            elif "前进" in response or "向前" in response:
                self.walk_forward()
            elif "后退" in response or "向后" in response:
                self.walk_backward()
            elif "趴下" in response or "躺下" in response:
                self.lie_down()
            elif "坐下" in response:
                self.sit()
            elif "停止" in response or "空闲" in response:
                self.idle()
            elif "跑" in response:
                self.run()
            elif "左转" in response or "向左" in response:
                self.turn_left()
            elif "右转" in response or "向右" in response:
                self.turn_right()
            return response
        except Exception as e:
            print(f"[Coze] 处理错误: {str(e)}")
            # 确保即使出错也恢复语音识别
            self.resume_voice_recognition()
            return None
            
    def close(self):
        """关闭所有连接"""
        print("✓ 关闭所有连接...")
        # 恢复语音识别确保能正常关闭
        self.resume_voice_recognition()
        self.voice_recognizer.stop()
        if self.spi:
            self.spi.close()
            print("✓ SPI连接已关闭")

def get_char(timeout=0.1):
    """非阻塞获取单个字符输入 - 服务环境安全版"""
    # 检查是否在终端中运行（关键修复！）
    if not sys.stdin.isatty():
        return None  # 服务环境中直接返回 None
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        tty.setcbreak(fd)
        rlist, _, _ = select.select([sys.stdin], [], [], timeout)
        if rlist:
            return sys.stdin.read(1)
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
    return None

def print_help():
    """打印帮助信息"""
    print("\n===== 机器人控制终端 =====")
    print("\n可用命令:")
    print("  1 - 站立")
    print("  2 - 向前走")
    print("  3 - 向后走")
    print("  4 - 趴下")
    print("  5 - 坐下")
    print("  6 - 奔跑")
    print("  7 - 向左转")
    print("  8 - 向右转")
    print("  0 - 停止/空闲")
    print("  s - 查询当前状态")
    print("  c - 与Coze AI对话")
    print("  h - 显示帮助")
    print("  q - 退出程序")
    print("\n语音命令:")
    print("  站立, 站起来, 站")
    print("  前进, 向前, 走")
    print("  后退, 向后")
    print("  奔跑, 跑")
    print("  向左转, 左转, 向左")
    print("  向右转, 右转, 向右")
    print("  趴下, 躺下, 趴")
    print("  坐下, 坐")
    print("  停止, 停, 空闲")
    print("  聊天/对话/问问 - 与Coze AI对话")

def main():
    robot = RobotController()
    print("\n===== 机器人控制终端 =====")
    print("当前状态:", robot.get_status())
    print("支持键盘、语音和Coze AI对话控制")
    print_help()
    
    # 主控制循环
    try:
        while True:
            # 检查语音命令
            voice_cmd = robot.voice_recognizer.get_command()
            if voice_cmd:
                # 执行语音命令
                if voice_cmd == "stand":
                    robot.stand()
                elif voice_cmd == "forward":
                    robot.walk_forward()
                elif voice_cmd == "backward":
                    robot.walk_backward()
                elif voice_cmd == "run":
                    robot.run()
                elif voice_cmd == "turn_left":
                    robot.turn_left()
                elif voice_cmd == "turn_right":
                    robot.turn_right()
                elif voice_cmd == "lie":
                    robot.lie_down()
                elif voice_cmd == "sit":
                    robot.sit()
                elif voice_cmd == "stop":
                    robot.idle()
                elif voice_cmd == "status":
                    print("当前状态:", robot.get_status())
                elif voice_cmd == "coze_start":
                    print("\n✓ 请说出你想问Coze的问题...")
                elif voice_cmd.startswith("coze:"):
                    # 处理Coze对话
                    text = voice_cmd[5:]
                    robot.process_coze_command(text)
                    
            # 非阻塞获取键盘输入
            key = get_char()
            if key:
                if key == '1':
                    robot.stand()
                elif key == '2':
                    robot.walk_forward()
                elif key == '3':
                    robot.walk_backward()
                elif key == '4':
                    robot.lie_down()
                elif key == '5':
                    robot.sit()
                elif key == '0':
                    robot.idle()
                elif key == 's':
                    print("当前状态:", robot.get_status())
                elif key == 'c':
                    print("\n请输入你想问Coze的问题: ", end="", flush=True)
                    user_input = input().strip()
                    if user_input:
                        robot.process_coze_command(user_input)
                elif key == 'h':
                    print_help()
                elif key in ['q', '\x03']:  # Ctrl+C or 'q'
                    print("\n✓ 退出程序...")
                    break
                else:
                    print(f"✗ 未知命令: {key} (输入'h'查看帮助)")
                    
            # 添加短暂延迟，降低CPU使用率
            time.sleep(0.01)
            
    except KeyboardInterrupt:
        print("\n✓ 程序被中断")
    finally:
        # 确保清理资源
        robot.close()
        print("✓ 程序已退出")

if __name__ == "__main__":   
    main()